# webapp.py
import gradio as gr
import numpy as np
import pickle
import cv2
import logging
import os

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class CatDogClassifier:
    def __init__(self):
        """初始化猫狗分类器"""
        self.model = None
        self.model_info = None
        self.model_loaded = False
        self.load_model()
        
    def load_model(self):
        """加载训练好的模型"""
        try:
            if os.path.exists("best_lazypredict_model.pkl"):
                with open("best_lazypredict_model.pkl", 'rb') as f:
                    self.model = pickle.load(f)
                
                # 加载模型信息
                if os.path.exists("model_info.pkl"):
                    with open("model_info.pkl", 'rb') as f:
                        self.model_info = pickle.load(f)
                
                self.model_loaded = True
                logging.info("模型加载成功")
                
                if self.model_info:
                    logging.info(f"模型类型: {self.model_info.get('model_name', '未知')}")
                    logging.info(f"准确率: {self.model_info.get('accuracy', 0):.4f}")
            else:
                logging.warning("模型文件不存在，请先运行 train.py 训练模型")
                self.model_loaded = False
        except Exception as e:
            logging.error(f"加载模型失败: {e}")
            self.model_loaded = False
    
    def preprocess_image(self, input_image):
        """
        预处理上传的图片
        """
        try:
            # Gradio返回的是numpy数组，直接处理
            if input_image is None:
                raise ValueError("图片数据为空")
            
            # 确保是numpy数组
            if isinstance(input_image, np.ndarray):
                img_array = input_image
            else:
                # 如果是文件路径，读取图片
                img_array = cv2.imread(input_image)
                if img_array is None:
                    raise ValueError("无法从路径读取图片")
            
            # 转换颜色空间并调整大小
            if len(img_array.shape) == 3:
                if img_array.shape[2] == 3:  # RGB图片
                    img_gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
                else:
                    img_gray = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)
            else:
                img_gray = img_array  # 已经是灰度图
            
            # 调整大小并展平
            img_resized = cv2.resize(img_gray, (32, 32))
            img_flatten = img_resized.flatten()
            img_normalized = img_flatten.astype('float32') / 255.0
            
            return img_normalized.reshape(1, -1)
                
        except Exception as e:
            logging.error(f"图片预处理失败: {e}")
            logging.error(f"输入图片类型: {type(input_image)}")
            if hasattr(input_image, 'shape'):
                logging.error(f"图片形状: {input_image.shape}")
            raise e
    
    def predict(self, input_image):
        """
        预测图片是猫还是狗
        """
        if not self.model_loaded:
            return {
                'prediction': '模型未加载',
                'confidence': 0.0,
                'probabilities': {'猫': 0.0, '狗': 0.0},
                'error': '请先运行 train.py 训练模型'
            }
        
        try:
            # 预处理图片
            processed_image = self.preprocess_image(input_image)
            
            # 进行预测
            if hasattr(self.model, 'predict_proba'):
                prediction = self.model.predict(processed_image)[0]
                probability = self.model.predict_proba(processed_image)[0]
            else:
                # 对于没有predict_proba的模型，使用predict
                prediction = self.model.predict(processed_image)[0]
                probability = [0.5, 0.5]  # 默认概率
            
            # 获取预测标签和置信度
            class_names = ['猫', '狗']
            predicted_class = class_names[prediction]
            confidence = probability[prediction] if len(probability) > prediction else 0.5
            
            # 返回结果
            result = {
                'prediction': predicted_class,
                'confidence': float(confidence),
                'probabilities': {
                    '猫': float(probability[0]) if len(probability) > 0 else 0.5,
                    '狗': float(probability[1]) if len(probability) > 1 else 0.5
                },
                'model_name': self.model_info.get('model_name', '未知模型') if self.model_info else '未知模型'
            }
            
            logging.info(f"预测结果: {result}")
            return result
            
        except Exception as e:
            logging.error(f"预测失败: {e}")
            return {
                'prediction': '错误',
                'confidence': 0.0,
                'probabilities': {'猫': 0.0, '狗': 0.0},
                'error': str(e)
            }

# 创建分类器实例
classifier = CatDogClassifier()

def classify_image(input_image):
    """
    Gradio接口函数，用于分类上传的图片
    """
    if not classifier.model_loaded:
        return "请先运行 train.py 训练模型", 0.5, 0.5, "未知模型"
    
    try:
        # 进行预测
        result = classifier.predict(input_image)
        
        if 'error' in result:
            return f"预测错误: {result['error']}", 0.5, 0.5, "未知模型"
        
        # 格式化输出
        prediction_text = f"预测结果: {result['prediction']} (置信度: {result['confidence']:.2%})"
        model_info = f"使用模型: {result.get('model_name', '未知模型')}"
        
        return prediction_text, result['probabilities']['猫'], result['probabilities']['狗'], model_info
        
    except Exception as e:
        error_msg = f"处理图片时发生错误: {str(e)}"
        logging.error(error_msg)
        return error_msg, 0.5, 0.5, "未知模型"

def create_interface():
    """创建Gradio界面"""
    
    description = """
    # 🐱🐶 猫狗分类器 (LazyPredict优化版)
    
    使用 LazyPredict 自动选择的最佳模型进行猫狗分类！
    
    **使用说明:**
    1. 首先运行 `train.py` 训练模型（只需要训练一次）
    2. 上传猫狗图片或使用示例图片
    3. 查看分类结果和概率
    
    **注意:** 请确保已运行 train.py 生成模型文件
    """
    
    with gr.Blocks(theme=gr.themes.Soft(), title="猫狗分类器 - LazyPredict版") as interface:
        gr.Markdown(description)
        
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(
                    label="上传图片",
                    type="numpy",
                    source="upload",
                    interactive=True
                )
                
            with gr.Column():
                result_output = gr.Textbox(
                    label="预测结果",
                    placeholder="这里将显示预测结果...",
                    interactive=False
                )
                
                model_info_output = gr.Textbox(
                    label="模型信息",
                    placeholder="这里将显示使用的模型信息...",
                    interactive=False
                )
                
                with gr.Row():
                    cat_prob = gr.Number(
                        label="猫的概率",
                        minimum=0,
                        maximum=1,
                        precision=4
                    )
                    
                    dog_prob = gr.Number(
                        label="狗的概率", 
                        minimum=0,
                        maximum=1,
                        precision=4
                    )
        
        # 绑定事件
        image_input.change(
            fn=classify_image,
            inputs=[image_input],
            outputs=[result_output, cat_prob, dog_prob, model_info_output]
        )
        
        # 添加示例图片
        gr.Examples(
            examples=[
                ["examples/cat_example.jpg"] if os.path.exists("examples/cat_example.jpg") else [],
                ["examples/dog_example.jpg"] if os.path.exists("examples/dog_example.jpg") else []
            ],
            inputs=image_input,
            label="示例图片 (请确保examples文件夹存在并包含图片)"
        )
    
    return interface

if __name__ == "__main__":
    interface = create_interface()
    logging.info("启动Gradio服务...")
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        inbrowser=True
    )